import pandas as pd
import random
from datetime import datetime, timedelta, timezone
import pymongo
from tqdm import tqdm

# ====================
# 初始化配置
# ====================
NUM_USERS = 20
NUM_REVIEWS = 200
FOCUS_RATIO = 0.75  # 90% 的评论集中在主要图书上
FOCUS_COUNT = 200    # 主要图书数量
OTHER_COUNT = 50   # 次要图书数量

# ====================
# 辅助函数
# ====================

def load_book_ids(csv_path):
    """从CSV加载图书ID"""
    try:
        df = pd.read_csv(csv_path)
        book_ids = df['_id'].tolist()
        print(f"从 {csv_path} 加载了 {len(book_ids)} 本图书的ID")
        return book_ids
    except Exception as e:
        print(f"加载图书ID时出错: {e}")
        return [f"book_{i}" for i in range(1, 101)]  # 默认生成100个图书ID

def generate_iso8601_time():
    """生成符合ISO 8601格式的时间（包含毫秒和时区）"""
    now = datetime.now(timezone.utc)
    random_seconds = random.randint(0, 30 * 24 * 60 * 60)  # 30天内的随机秒数
    random_microseconds = random.randint(0, 999999)
    dt = now - timedelta(seconds=random_seconds)
    return dt.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00"

def generate_review_content(rating):
    """根据评分生成评论内容"""
    if rating >= 4:
        return random.choice([
            "这本书太棒了，强烈推荐！",
            "作者写得非常好，我一口气读完了",
            "故事情节引人入胜，值得一读再读"
        ])
    elif rating == 3:
        return random.choice([
            "整体还不错，但有些部分可以改进",
            "中规中矩的一本书",
            "没有特别惊艳，但也不差"
        ])
    else:
        return random.choice([
            "不太推荐这本书，",
            "读起来有些失望，",
            "感觉内容有些浅薄，"
        ])

def save_to_csv(df, output_path):
    """保存为CSV文件"""
    df.to_csv(output_path, index=False, encoding='utf-8-sig')
    print(f"数据已保存到 {output_path} (共{len(df)}条记录)")

def save_to_mongodb(df, db_name='test', collection_name='review'):
    """保存到MongoDB"""
    client = pymongo.MongoClient("mongodb://localhost:27017/")
    db = client[db_name]
    collection = db[collection_name]
    collection.delete_many({})
    # 转换时间字符串为datetime对象
    records = df.to_dict('records')
    for record in records:
        try:
            record['createTime'] = datetime.strptime(record['createTime'], "%Y-%m-%dT%H:%M:%S.%f%z")
            record['updateTime'] = datetime.strptime(record['updateTime'], "%Y-%m-%dT%H:%M:%S.%f%z")
        except ValueError:
            # 如果时间格式不匹配，跳过转换
            pass

    result = collection.insert_many(records)
    print(f"成功插入 {len(result.inserted_ids)} 条数据到MongoDB {db_name}.{collection_name}")

def load_or_generate_books(csv_path):
    """加载图书ID或生成默认值"""
    try:
        book_ids = load_book_ids(csv_path)
        if len(book_ids) < FOCUS_COUNT + OTHER_COUNT:
            raise ValueError("图书ID数量不足，需至少包含主要图书和次要图书")
        return book_ids
    except Exception as e:
        print(f"加载图书ID失败，使用默认值: {e}")
        return [f"book_{i}" for i in range(1, FOCUS_COUNT + OTHER_COUNT + 1)]

def select_focus_and_other_books(book_ids):
    """选择主要图书和次要图书"""
    focus_books = random.sample(book_ids, min(FOCUS_COUNT, len(book_ids)))
    other_books = random.sample([book for book in book_ids if book not in focus_books], min(OTHER_COUNT, len(book_ids) - len(focus_books)))
    return focus_books, other_books

def generate_pending_reviews(book_ids, num_users=NUM_USERS, num_reviews=NUM_REVIEWS):
    """
    生成待审核书评数据ID更集中
    """
    # 分类图书
    focus_books, other_books = select_focus_and_other_books(book_ids)
    print(f"主要图书ID: {focus_books}")
    print(f"次要图书ID: {other_books}")

    # 计算评论分配
    total_focus = int(num_reviews * FOCUS_RATIO)
    total_other = num_reviews - total_focus

    print(f"生成 {total_focus} 条主要图书评论，{total_other} 条次要图书评论。")

    reviews = []

    # 生成主要图书评论
    for _ in tqdm(range(total_focus), desc="生成主要图书评论"):
        create_time = generate_iso8601_time()
        review = {
            "u_id": random.randint(1, NUM_USERS),
            "b_id": random.choice(focus_books),
            "content": generate_review_content(random.randint(1, 5)),
            "rating": random.randint(1, 5),
            "status": 0,  # 待审核
            "createTime": create_time,
            "updateTime": create_time
        }
        reviews.append(review)

    # 生成次要图书评论
    for _ in tqdm(range(total_other), desc="生成次要图书评论"):
        create_time = generate_iso8601_time()
        review = {
            "u_id": random.randint(1, NUM_USERS),
            "b_id": random.choice(other_books),
            "content": generate_review_content(random.randint(1, 5)),
            "rating": random.randint(1, 5),
            "status": 0,  # 待审核
            "createTime": create_time,
            "updateTime": create_time
        }
        reviews.append(review)

    # 转换为DataFrame
    reviews_df = pd.DataFrame(reviews)

    # 统计图书ID分布
    book_distribution = reviews_df['b_id'].value_counts()
    print("图书ID分布情况：")
    print(book_distribution)

    return reviews_df

if __name__ == "__main__":
    # 1. 加载或生成图书ID
    book_ids = load_or_generate_books("../../../file/ids.csv")  # 替换为你的CSV文件路径

    # 2. 生成待审核书评数据
    reviews_df = generate_pending_reviews(
        book_ids=book_ids,
        num_users=NUM_USERS,
        num_reviews=NUM_REVIEWS
    )

    # 3. 保存数据
    save_to_csv(reviews_df, "pending_reviews_iso.csv")
    save_to_mongodb(reviews_df)